import cv2
import numpy as np
from ultralytics import YOLO

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT,720)

model = YOLO('yolov5su.pt')

while True:
        ret, frame = cap.read()

        if not ret:
                print("exit")
                break

        res = model(frame)

        # cls 检测到的物体类型, 是一个列表, 包含图片中所有检测到的物体类型
        # 比如说检测到人, 杯子, 键盘, 这 cls 列表有三个元素, 分辨表示 人, 杯子, 键盘的枚举值
        cls = res[0].boxes.cls

        # xyxy 检测到的物体为止, 也是一个列表, 包含图片中所有检测到的物体位置
        xyxy = res[0].boxes.xyxy

        p1 = (0, 0)
        p2 = (0, 0)
        detect = False

        # 遍历所有的检测结果, 判断是否有检测到人
        for i, t in enumerate(cls):
                if t == 0:                      # 0 表示 person 类型
                        p1 = (int(xyxy[i][0].item()), int(xyxy[i][1].item()))
                        p2 = (int(xyxy[i][2].item()), int(xyxy[i][3].item()))
                        print(xyxy[i])
                        print(p1)
                        print(p2)
                        detect = True
                        break

        # 如果有检测到人, 则叠加矩形框
        color = (0, 255, 255)
        thickness = 2
        if detect:
                frame = cv2.rectangle(frame, p1, p2, color, thickness)

        cv2.imshow('camere', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
                break

cap.release()
cv2.destroyAllWindows()

